package com.ai.module.mj.loadbalancer;

import com.ai.common.BaseResp;
import com.ai.module.business.service.BusinessTotalService;
import com.ai.module.business.service.ImageGroundService;
import com.ai.module.chat.common.base.ThreadLocalContext;
import com.ai.module.chat.common.constant.StringConstant;
import com.ai.module.chat.config.advice.exception.BusinessException;
import com.ai.module.chat.config.properties.AliyunProperty;
import com.ai.module.chat.util.DateUtils;
import com.ai.module.chat.util.RedisTemplateUtils;
import com.ai.module.chat.util.SpringUtils;
import com.ai.module.mj.constant.enums.BlendDimensions;
import com.ai.module.mj.constant.enums.TaskStatus;
import com.ai.module.mj.constant.stringConstant.Constants;
import com.ai.module.mj.constant.stringConstant.ReturnCode;
import com.ai.module.mj.domain.DiscordAccount;
import com.ai.module.mj.domain.Task;
import com.ai.module.mj.domain.bo.MJDataUrl;
import com.ai.module.mj.domain.vo.SubmitResultVO;
import com.ai.module.mj.service.*;
import com.ai.module.mj.service.impl.DiscordServiceImpl;
import com.ai.module.mj.wss.UserWebSocketStarter;
import com.ai.module.mj.wss.WebSocketStarter;
import com.ai.module.user.domain.entity.UserTaskEntity;
import com.ai.module.user.domain.vo.UserLoginInfo;
import com.ai.module.user.service.UserTaskService;
import com.ai.module.websocket.service.AIProductWebSocketService;
import com.ai.module.websocket.service.GlobalWebSocketService;
import com.ai.module.websocket.vo.WebSocketProduct;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.socket.WebSocketSession;

import java.util.*;
import java.util.concurrent.Callable;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;

@Slf4j
public class DiscordInstanceImpl implements DiscordInstance {

    public static final  String TASK_REDIS_KEY= "user_Task:";
    private final DiscordAccount account;
    private final WebSocketStarter socketStarter;
    private final DiscordService service;
    private final HighTaskStoreService highTaskStoreService;
    // 低优先级队列 存储
    private final LowTaskStoreService lowTaskStoreService= SpringUtils.getClazz(LowTaskStoreService.class);

    private final NotifyService notifyService;

    // 高优先级队列 todo 低优先级队列
    private final ThreadPoolTaskExecutor highTaskExecutor;

    // 高优先级队列 todo 低优先级队列
    private final ThreadPoolTaskExecutor slowTaskExecutor;

    private final List<Task> runningTasks;
    private final Map<String, Future<?>> taskFutureMap = Collections.synchronizedMap(new HashMap<>());

    private ImageGroundService imageGroundService = SpringUtils.getClazz(ImageGroundService.class);
    private UserTaskService taskService =  SpringUtils.getClazz(UserTaskService.class);


    private AliyunProperty aliyunProperty = SpringUtils.getClazz(AliyunProperty.class);

    private BusinessTotalService businessTotalService= SpringUtils.getClazz(BusinessTotalService.class);

    public DiscordInstanceImpl(DiscordAccount account, UserWebSocketStarter socketStarter, RestTemplate restTemplate,
                               HighTaskStoreService highTaskStoreService, NotifyService notifyService, Map<String, String> paramsMap) {
        this.account = account;
        this.socketStarter = socketStarter;
        this.highTaskStoreService = highTaskStoreService;
        this.notifyService = notifyService;
        this.service = new DiscordServiceImpl(account, restTemplate, paramsMap);
        this.runningTasks = new CopyOnWriteArrayList<>();
        this.highTaskExecutor = new ThreadPoolTaskExecutor();
        this.highTaskExecutor.setCorePoolSize(account.getHighCoreSize());
        this.highTaskExecutor.setMaxPoolSize(account.getHighCoreSize());
        this.highTaskExecutor.setQueueCapacity(account.getHighQueueSize());
        this.highTaskExecutor.setThreadNamePrefix("high-TaskQueue-" + account.getDisplay() + "-");
        this.highTaskExecutor.initialize();
        // 低优先级线程池初始化 todo 移到 beanConfig
        this.slowTaskExecutor = new ThreadPoolTaskExecutor();
        this.slowTaskExecutor.setCorePoolSize(account.getSlowCoreSize());
        this.slowTaskExecutor.setMaxPoolSize(account.getSlowCoreSize());
        this.slowTaskExecutor.setQueueCapacity(account.getSlowQueueSize());
        this.slowTaskExecutor.setThreadNamePrefix("slow-TaskQueue-" + account.getDisplay() + "-");
        this.slowTaskExecutor.initialize();

        this.lowTaskStoreService.setMaxQue(account.getHighQueueSize());
        this.highTaskStoreService.setMaxQue(account.getHighQueueSize());

    }

    @Override
    public String getInstanceId() {
        return this.account.getChannelId();
    }

    @Override
    public DiscordAccount account() {
        return this.account;
    }

    @Override
    public boolean isAlive() {
        return this.account.isEnable();
    }

    @Override
    public void startWss() throws Exception {
        this.socketStarter.setTrying(true);
        this.socketStarter.start();
    }

    @Override
    public List<Task> getRunningTasks() {
        return this.runningTasks;
    }

    @Override
    public void exitTask(Task task) {
        try {
            Future<?> future = this.taskFutureMap.get(task.getId());
            if (future != null) {
                future.cancel(true);
            }
            updateStatus(task);
        } finally {
            log.error("开始任务退出机制 [{}]", task.getId());
            this.runningTasks.remove(task);
            this.taskFutureMap.remove(task.getId());
        }
    }

    @Override
    public Map<String, Future<?>> getRunningFutures() {
        return this.taskFutureMap;
    }

    @Override
    public synchronized SubmitResultVO submitTask(Task task, Callable<BaseResp<Void>> discordSubmit) {
        task.setUserId(ThreadLocalContext.getUserId());
        var userId =Optional.ofNullable(ThreadLocalContext.getUserLoginInfo())
                .map(UserLoginInfo::getId).orElse(0);
        if (userId ==0 ){
            log.error("出现了用户信息ID获取失败情况 userInfo [{}]",ThreadLocalContext.getUserLoginInfo());
        }
        task.setUserIntId(userId);

        TaskStoreService taskStoreService = highTaskStoreService;
        ThreadPoolTaskExecutor threadPoolTaskExecutor = highTaskExecutor;
        if ( !Objects.equals(task.getDrawType(), "FAST")){
            taskStoreService = lowTaskStoreService;
            threadPoolTaskExecutor= slowTaskExecutor;
        }
        taskStoreService.save(task);
        int currentWaitNumbers;
        try {
            currentWaitNumbers = threadPoolTaskExecutor.getThreadPoolExecutor().getQueue().size();
            Future<?> future =   threadPoolTaskExecutor.submit(() -> executeTask(task, discordSubmit));
            this.taskFutureMap.put(task.getId(), future);
        } catch (RejectedExecutionException e) {
            taskStoreService.delete(task.getId());
            UserTaskEntity userTask = new UserTaskEntity();
            BeanUtils.copyProperties(task,userTask);
            userTask.setPrompt("队列已满，请稍后尝试");
            taskService.save(userTask);
            throw new BusinessException(ReturnCode.QUEUE_REJECTED, "队列已满，请稍后尝试") ;
        } catch (Exception e) {
            UserTaskEntity userTask = new UserTaskEntity();
            BeanUtils.copyProperties(task,userTask);
            userTask.setPrompt("提交失败，系统异常");
            taskService.save(userTask);
            log.error("submit task error", e);
            e.printStackTrace();
            throw new BusinessException(ReturnCode.FAILURE, "提交失败，系统异常");

            //	.setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, this.getInstanceId());
        }
        // 进行任务保存
        UserTaskEntity userTask = new UserTaskEntity();
        BeanUtils.copyProperties(task,userTask);
        userTask.setTaskId(task.getId());
        taskService.save(userTask);
        RedisTemplateUtils.set(TASK_REDIS_KEY,userTask.getTaskId(), DateUtils.TIME_DAY  );
        if (currentWaitNumbers == 0) {
            String desc = "正在绘画中,请稍等....";
            sendNoEndImage(task,currentWaitNumbers);
            return SubmitResultVO.of(ReturnCode.SUCCESS, desc, task.getId())
                    .setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, this.getInstanceId());
        } else {
            // 进行 任务保存
            String desc = "排队中，前面还有" + currentWaitNumbers + "个任务";
            sendNoEndImage(task,currentWaitNumbers);
            log.info("正在排队中,前面还有[{}]个任务",currentWaitNumbers);
            return SubmitResultVO.of(ReturnCode.IN_QUEUE, desc, task.getId())
                    .setProperty("numberOfQueues", currentWaitNumbers)
                    .setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, this.getInstanceId());
        }
    }

    private void executeTask(Task task, Callable<BaseResp<Void>> discordSubmit) {
        this.runningTasks.add(task);
        String taskId = task.getId();
        try {
            task.start();
            BaseResp<Void> result = discordSubmit.call();
            if (result.getCode() != ReturnCode.SUCCESS) {
                task.fail(result.getMessage());
                updateStatus(task);
                return;
            }
            //
            updateStatus(task);
            do {
                if (! RedisTemplateUtils.hasKey(TASK_REDIS_KEY+taskId )) {
                    UserTaskEntity userTask = new UserTaskEntity();
                    userTask.setStatus(task.getStatus());
                    taskService.update(userTask,new LambdaQueryWrapper<UserTaskEntity>()
                            .eq(UserTaskEntity::getTaskId, task.getUserId()));
                    RedisTemplateUtils.set(TASK_REDIS_KEY+taskId,task, DateUtils.TIME_DAY);
                }
                task.sleep(3000L);
                updateStatus(task);
            } while (task.getStatus() == TaskStatus.IN_PROGRESS
                    || task.getStatus() == TaskStatus.SUBMITTED);
            log.info("task finished, id: {}, status: {}", taskId, task.getStatus());
            // 画完了  这个时候可以通知前端,删除任务之前可以考虑将数据存到 redis
        } catch (InterruptedException e) {
            e.printStackTrace();
            log.error("图片处理异常啊 message [{}]", e.getMessage());
            task.fail("执行错误，系统异常");
            updateStatus(task);
            Thread.currentThread().interrupt();

        } catch (Exception e) {
            e.printStackTrace();
            log.error("task execute error", e);
            task.fail("执行错误，系统异常");
            updateStatus(task);
            task.awake();
        } finally {
            // todo 开始退出任务
            updateStatus(task);
            this.runningTasks.remove(task);
            this.taskFutureMap.remove(taskId);
            log.error("开始退出任务啊 [{}], 还剩下任务数 [{}]",taskId,runningTasks.size());

        }
    }




    private void updateStatus(Task task) {
        if (task.getStatus()!= TaskStatus.SUCCESS){
            task.setStatus(TaskStatus.IN_PROGRESS);
            this.highTaskStoreService.save(task);
        }
        this.notifyService.notifyTaskChange(task);
    }

    @Override
    public BaseResp<Void> imagine(String prompt, String nonce) {
        return this.service.imagine(prompt, nonce);
    }

    @Override
    public BaseResp<Void> upscale(String messageId, int index, String messageHash, int messageFlags, String nonce) {
        return this.service.upscale(messageId, index, messageHash, messageFlags, nonce);
    }

    @Override
    public BaseResp<Void> variation(String messageId, int index, String messageHash, int messageFlags, String nonce) {
        return this.service.variation(messageId, index, messageHash, messageFlags, nonce);
    }

    @Override
    public BaseResp<Void> reroll(String messageId, String messageHash, int messageFlags, String nonce) {
        return this.service.reroll(messageId, messageHash, messageFlags, nonce);
    }

    @Override
    public BaseResp<Void> describe(String finalFileName, String nonce) {
        return this.service.describe(finalFileName, nonce);
    }

    @Override
    public BaseResp<Void> blend(List<String> finalFileNames, BlendDimensions dimensions, String nonce) {
        return this.service.blend(finalFileNames, dimensions, nonce);
    }

    @Override
    public BaseResp<String> upload(String fileName, MJDataUrl dataUrl) {
        return this.service.upload(fileName, dataUrl);
    }

    @Override
    public BaseResp<String> sendImageMessage(String content, String finalFileName) {
        return this.service.sendImageMessage(content, finalFileName);
    }

    public void sendNoEndImage(Task task, int waitNum) {
        String userId = task.getUserId();
        String robotId = task.getRobotId();
        String userRobotKey = String.format(StringConstant.userROBOT, userId, robotId);
        Map<String, WebSocketSession> userProductSessionMap = GlobalWebSocketService.userProductSessionMap;
        if (userProductSessionMap.containsKey(userRobotKey)) {
            WebSocketSession webSocketSession = userProductSessionMap.get(userRobotKey);
            WebSocketProduct product = new WebSocketProduct();
            WebSocketProduct.AiImage aiImage = new WebSocketProduct.AiImage();
            aiImage.setProcess(task.getProgress());
//			aiImage.setTask(task);
            aiImage.setTaskId(task.getId());
            product.setAiImage(aiImage);
            if (waitNum!= 0){
                aiImage.setWaiNumber(waitNum);
            }
            SpringUtils.getClazz(AIProductWebSocketService.class)  .sendMessage(webSocketSession,product);
        }
    }

}
